package main

import (
	"context"
	"fmt"
	"time"
)

// RateLimiter 接口限流结构体
type RateLimiter struct {
	rate     int           // 限制的速率
	interval time.Duration // 单位时间间隔
	tokens   chan struct{} // 令牌的channel
	stop     chan struct{} // 停止信号的channel
}

// NewRateLimiter 初始化该结构体
func NewRateLimiter(rate int, interval time.Duration) *RateLimiter {
	limiter := &RateLimiter{
		rate:     rate,
		interval: interval,
		tokens:   make(chan struct{}, rate), //token的channel缓冲区
		stop:     make(chan struct{}),
	}
	go limiter.start()
	return limiter
}

// start 启动定时器
func (r *RateLimiter) start() {
	ticker := time.NewTicker(r.interval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			for i := 0; i < r.rate; i++ {
				select {
				case r.tokens <- struct{}{}:
				default:
					// 如果tokens缓冲区已满，则不再发送token
				}
			}
		case <-r.stop:
			close(r.tokens) //关闭通道
			return
		}
	}
}

// Allow 是否允许继续执行 否则进行阻塞
func (r *RateLimiter) Allow(ctx context.Context) bool {
	select {
	case <-r.tokens:
		return true
	case <-ctx.Done():
		return false
	}
}

// 通知定时器任务
func (r *RateLimiter) Stop() {
	r.stop <- struct{}{}
}

func main() {
	limiter := NewRateLimiter(3, time.Second) // 每秒限制3个请求
	defer limiter.Stop()

	for i := 1; i <= 10; i++ {
		go execute(limiter, i)
	}
	time.Sleep(time.Second * 10)
}

// execute 执行任务
func execute(limiter *RateLimiter, i int) {
	fmt.Println("协程:", time.Now(), i)
	ctx, cf := context.WithTimeout(context.Background(), time.Second*3)
	defer cf()
	if limiter.Allow(ctx) {
		fmt.Println("Request", time.Now(), i, "allowed")
	} else {
		fmt.Println("Request", time.Now(), i, "timeout")
	}
}
